import paddle
import paddle.nn as nn
import paddle.optimizer as optim
import numpy as np
from model_define import FallDetectionLSTM

# === 1. 构建模拟数据集 ===
# 假设输入特征：[心率, 呼吸率, 体动幅度, 信号强度]
def generate_dummy_data(samples=1000, time_steps=30):
    X = []
    y = []
    for _ in range(samples):
        # 50% 概率模拟跌倒
        if np.random.rand() > 0.5:
            # 跌倒数据：特征值波动大，突变
            data = np.random.normal(loc=0.8, scale=0.2, size=(time_steps, 4))
            label = 1
        else:
            # 正常数据：特征值平稳
            data = np.random.normal(loc=0.2, scale=0.1, size=(time_steps, 4))
            label = 0
        X.append(data)
        y.append(label)
    
    # 转换为 Paddle Tensor
    return paddle.to_tensor(np.array(X), dtype='float32'), \
           paddle.to_tensor(np.array(y), dtype='int64')

# === 2. 训练流程 ===
def train():
    print("正在生成训练数据...")
    train_X, train_y = generate_dummy_data()
    
    model = FallDetectionLSTM()
    model.train()
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(learning_rate=0.001, parameters=model.parameters())
    
    epochs = 10
    print(f"开始训练，共 {epochs} 轮...")
    
    for epoch in range(epochs):
        # 前向计算
        logits = model(train_X)
        loss = criterion(logits, train_y)
        
        # 反向传播
        loss.backward()
        optimizer.step()
        optimizer.clear_grad()
        
        # 打印日志
        if epoch % 2 == 0:
            acc = paddle.metric.accuracy(input=logits, label=train_y.unsqueeze(1))
            print(f"Epoch {epoch}: Loss = {loss.numpy()[0]:.4f}, Acc = {acc.numpy()[0]:.4f}")

    # === 3. 保存模型 ===
    paddle.save(model.state_dict(), "fall_detection_model.pdparams")
    print("模型已保存为 fall_detection_model.pdparams")

if __name__ == "__main__":
    train()